from skimage.transform import rotate
from skimage.feature import local_binary_pattern
from skimage import data, io, data_dir, filters, feature
from skimage.color import label2rgb
import skimage
import torch
import numpy as np
import matplotlib.pyplot as plt
import os
from PIL import Image
import cv2


def show_format_image(image, figure=0):
    HSV = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
    YCrCb = cv2.cvtColor(image, cv2.COLOR_BGR2YCrCb)
    LAB = cv2.cvtColor(image, cv2.COLOR_BGR2LAB)
    plt.figure(figure)
    plt.subplot(221), plt.imshow(HSV), plt.title('HSV')
    plt.subplot(222), plt.imshow(cv2.split(HSV)[0], plt.cm.gray), plt.title('H')
    plt.subplot(223), plt.imshow(cv2.split(HSV)[1], plt.cm.gray), plt.title('S')
    plt.subplot(224), plt.imshow(cv2.split(HSV)[2], plt.cm.gray), plt.title('V')
    plt.figure(figure + 1)
    plt.subplot(221), plt.imshow(YCrCb), plt.title('YCrCb')
    plt.subplot(222), plt.imshow(cv2.split(YCrCb)[0], plt.cm.gray), plt.title('Y')
    plt.subplot(223), plt.imshow(cv2.split(YCrCb)[1], plt.cm.gray), plt.title('Cr')
    plt.subplot(224), plt.imshow(cv2.split(YCrCb)[2], plt.cm.gray), plt.title('Cb')
    # Cr = cv2.split(YCrCb)[1]
    # Cb = cv2.split(YCrCb)[2]
    # plt.figure(figure+3)
    # Y = np.ones((256,128),dtype=np.uint8)*127
    # Trans = cv2.merge([Y,Cr,Cb])
    # Trans = cv2.cvtColor(Trans,cv2.COLOR_YCrCb2RGB)
    # plt.imshow(Trans)
    # plt.show()
    plt.figure(figure + 2)
    plt.subplot(221), plt.imshow(LAB), plt.title('LAB')
    plt.subplot(222), plt.imshow(cv2.split(LAB)[0], plt.cm.gray), plt.title('L')
    plt.subplot(223), plt.imshow(cv2.split(LAB)[1], plt.cm.gray), plt.title('A')
    plt.subplot(224), plt.imshow(cv2.split(LAB)[2], plt.cm.gray), plt.title('B')


def show_YCrCb(image, figure=0):
    YCrCb = cv2.cvtColor(image, cv2.COLOR_BGR2YCrCb)
    RGB = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    plt.figure(figure)
    plt.subplot(221), plt.imshow(RGB), plt.title('YCrCb')
    plt.subplot(222), plt.imshow(cv2.split(YCrCb)[0], plt.cm.gray), plt.title('Y')
    plt.subplot(223), plt.imshow(cv2.split(YCrCb)[1], plt.cm.gray), plt.title('Cr')
    plt.subplot(224), plt.imshow(cv2.split(YCrCb)[2], plt.cm.gray), plt.title('Cb')


def show_HSV(image, figure=0):
    HSV = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
    RGB = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    plt.figure(figure)
    plt.subplot(221), plt.imshow(RGB), plt.title('HSV')
    plt.subplot(222), plt.imshow(cv2.split(HSV)[0], plt.cm.gray), plt.title('H')
    plt.subplot(223), plt.imshow(cv2.split(HSV)[1], plt.cm.gray), plt.title('S')
    plt.subplot(224), plt.imshow(cv2.split(HSV)[2], plt.cm.gray), plt.title('V')


def show_RGB(image, figure=0):
    RGB = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    plt.figure(figure)
    plt.subplot(221), plt.imshow(RGB), plt.title('HSV')
    plt.subplot(222), plt.imshow(cv2.split(RGB)[0], plt.cm.gray), plt.title('R')
    plt.subplot(223), plt.imshow(cv2.split(RGB)[1], plt.cm.gray), plt.title('G')
    plt.subplot(224), plt.imshow(cv2.split(RGB)[2], plt.cm.gray), plt.title('B')


def color_balance(img, figure=0):
    RGB = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    b, g, r = cv2.split(img)
    B = np.mean(b)
    G = np.mean(g)
    R = np.mean(r)
    K = (R + G + B) / 3
    Kb = K / B
    Kg = K / G
    Kr = K / R
    cv2.addWeighted(b, Kb, 0, 0, 0, b)
    cv2.addWeighted(g, Kg, 0, 0, 0, g)
    cv2.addWeighted(r, Kr, 0, 0, 0, r)
    merged = cv2.merge([b, g, r])
    plt.figure(figure)
    plt.subplot(121), plt.imshow(RGB), plt.title('S')
    plt.subplot(122), plt.imshow(cv2.cvtColor(merged, cv2.COLOR_BGR2RGB)), plt.title('T')


def show_gray(image, figure=0):
    RGB = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    Gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    cv2.equalizeHist(Gray, Gray)

    # kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
    # Gray = cv2.erode(Gray, kernel)  # 腐蚀图像
    # Gray = cv2.dilate(Gray,kernel)      #膨胀图像
    plt.figure(figure)
    plt.subplot(121), plt.imshow(RGB), plt.title('S')
    plt.subplot(122), plt.imshow(Gray, plt.cm.gray), plt.title('T')
def show_RGB_hist(img_bgr,figure=0):
    # 按R、G、B三个通道分别计算颜色直方图
    b_hist = cv2.calcHist([img_bgr], [0], None, [256], [0, 256])
    g_hist = cv2.calcHist([img_bgr], [1], None, [256], [0, 256])
    r_hist = cv2.calcHist([img_bgr], [2], None, [256], [0, 256])

    # 显示3个通道的颜色直方图
    plt.figure(figure)
    plt.plot(b_hist, label='B', color='blue')
    plt.plot(g_hist, label='G', color='green')
    plt.plot(r_hist, label='R', color='red')
    plt.legend(loc='best')
    plt.xlim([0, 256])
    # plt.show()

import copy
def hisEqulColor2(img, figure=0):
    RGB = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    # 将RGB图像转换到YCrCb空间中
    # ycrcb = cv2.cvtColor(img, cv2.COLOR_BGR2YCrCb)
    ycrcb = copy.deepcopy(RGB)
    # 将YCrCb图像通道分离
    channels = cv2.split(ycrcb)
    # 对第1个通道即亮度通道进行全局直方图均衡化并保存
    cv2.equalizeHist(channels[0], channels[0])
    cv2.equalizeHist(channels[1], channels[1])
    # channels[1]+=10
    cv2.equalizeHist(channels[2], channels[2])
    # 将处理后的通道和没有处理的两个通道合并，命名为ycrcb
    cv2.merge(channels, ycrcb)
    # 将YCrCb图像转换回RGB图像
    # cv2.cvtColor(ycrcb, cv2.COLOR_YCR_CB2RGB, img)
    # img = cv2.GaussianBlur(img,(3,3),0)
    img = ycrcb
    plt.figure(figure)
    plt.subplot(121), plt.imshow(RGB), plt.title('S')
    plt.subplot(122), plt.imshow(img, ), plt.title('T')
    return img


if __name__ == '__main__':
    # 读取图像
    image1 = cv2.imread('00000002.png')
    image2 = cv2.imread('00000008.png')
    image3 = cv2.imread('00000899.png')
    image4 = cv2.imread('00001996.png')
    show_RGB_hist(image1,0)
    show_RGB_hist(image2,1)
    show_RGB_hist(image3,2)
    show_RGB_hist(image4,3)
    plt.show()
    # show_format_image(image1, 0)
    # show_format_image(image2, 4)
    # show_format_image(image3, 8)
    # show_format_image(image4, 12)
    root_path = './red_blue'
    root_path = './green'
    root_path = './violet'
    # root_path = './white'
    # root_path = './yellow'
    image_list = os.listdir(root_path)
    for idx, image_name in enumerate(image_list):
        image = cv2.imread(root_path + '/' + image_name)
        # hisEqulColor2(image, idx)
        # show_YCrCb(image, idx)
        # show_HSV(image,idx)
        # show_gray(image,idx)
        # show_RGB(image,idx)
        show_RGB_hist(image,idx)
        # color_balance(image, idx)
    plt.show()
